#include "btreestore.h"
#include <pthread.h>
#define int_max 0xFFFFFFFF
pthread_mutex_t lock_mutex = PTHREAD_MUTEX_INITIALIZER;
uint8_t num_processors;
void * init_store(uint16_t branching, uint8_t n_processors) {
    num_processors = n_processors;
    struct b_tree *b_tree = malloc(sizeof(struct b_tree));
    b_tree->min_val = (branching - 1) / 2;
    b_tree->max_val = branching - 1;
    b_tree->mid_index = branching / 2;
    b_tree->root = NULL;
    return b_tree;
}

struct node *create_tree_node(struct b_tree *b_tree) {
    struct node * node = (struct node *) malloc(sizeof(struct node));

    node->num_keys = 0;

    node->keys = (uint32_t *) malloc(sizeof(uint32_t) * (b_tree->max_val + 1));
    memset(node->keys, 0, sizeof(uint32_t) * (b_tree->max_val + 1));

    node->data = (struct info *) malloc(sizeof(struct info) * (b_tree->max_val + 1));
    memset(node->data, 0, sizeof(struct info) * (b_tree->max_val + 1));

    node->child = (struct node **) malloc(sizeof(struct node *) * (b_tree->max_val + 2));
    memset(node->child, 0, sizeof(struct node *) * (b_tree->max_val + 2));

    node->parent = NULL;

    return node;

}

void* create_info(void * plaintext, size_t count, uint32_t encryption_key[4], uint64_t nonce, void *helper){
    struct b_tree *b_tree = helper;
    //encrypt plaintext
    uint32_t num_blocks = count % 8 == 0 ? count / 8 : (count / 8 + 1);
    uint64_t *plain = (uint64_t *)malloc(sizeof (uint64_t) * num_blocks);
    memset(plain, 0, sizeof (uint64_t) * num_blocks);
    uint64_t *cipher = (uint64_t *) malloc(sizeof (uint64_t) * num_blocks);
    memset(cipher, 0, sizeof (uint64_t) * num_blocks);

    memcpy(plain, plaintext, count);
    encrypt_tea_ctr(plain, encryption_key, nonce, cipher, num_blocks);

    if(plain){
        free(plain);
        plain = NULL;
    }
    return cipher;
}

void close_node(struct node* node, struct node* root){
    if(node != NULL){
        if(node->child != NULL && node->child[0]){
            for (uint16_t i = 0; i <= node->num_keys; i++) {
                close_node(node->child[i], root);
            }
        }
        if(node != root){
            free(node);
            node = NULL;
        }
    }
}

void close_store(void *helper) {
    struct b_tree * tree = helper;
    int sum = 0;
    count_tree_node(tree->root, &sum);
    struct node ** list = malloc(sizeof(struct node *) * sum);
    btree_export(helper, list);
    tree->root = NULL;
    free(helper);
    helper = NULL;
    for(int i = 0; i < sum; i++){
        struct node *p = list[i];
        list[i] = NULL;
        free(p);
    }
    free(list);

}

int btree_insert(uint32_t key, void *plaintext, size_t count, uint32_t encryption_key[4], uint64_t nonce,  void * helper) {
    pthread_mutex_lock(&lock_mutex);
    struct b_tree *b_tree = helper;
    struct node *node = b_tree->root;
    //empty tree
    if (node == NULL) {
        node = create_tree_node(b_tree);
        node->num_keys++;
        node->keys[0] = key;
        b_tree->root = node;
        node->parent = NULL;
        //init data
       // uint32_t num_blocks = count % 8 == 0 ? count / 8 : (count / 8 + 1);
        memcpy(node->data->key, encryption_key, sizeof (uint32_t) * 4);
        node->data->size = count;
        node->data->nonce = nonce;
        node->data->data = create_info(plaintext, count, encryption_key, nonce, helper);
        pthread_mutex_unlock(&lock_mutex);
        return 0;
    }
    //find the index of tree node key to insert
    uint16_t index = 0;
    while (node != NULL) {
        for (index = 0; index < node->num_keys; index++) {
            //key already exist in b tree
            if (node->keys[index] == key) {
                pthread_mutex_unlock(&lock_mutex);
                return 1;
            }
            if (node->keys[index] > key) {
                break;
            }
        }
        //search  child
        if (node->child[index] != NULL) {
            node = node->child[index];
        } else {
            break;
        }
    }
    for (uint16_t i = node->num_keys; i > index; i--) {
        node->keys[i] = node->keys[i - 1];
        node->data[i] = node->data[i - 1];
    }
    //insert into b tree
    node->num_keys++;
    node->keys[index] = key;
    (node->data + index)->size = count;
    (node->data + index)->nonce = nonce;
    memcpy((node->data + index)->key, encryption_key, sizeof (uint32_t) * 4);
    (node->data + index)->data = create_info(plaintext, count, encryption_key, nonce, helper);

    //overflow spilt node
    uint16_t mid_index = b_tree->mid_index;
    struct node *node_parent = NULL;
    struct node *right_node = NULL;
    while (b_tree->max_val < node->num_keys) {
        uint16_t num_count = node->num_keys;
        //init a node
        right_node = create_tree_node(b_tree);
        //copy array data
        memcpy(right_node->keys, node->keys + mid_index + 1, (num_count - mid_index - 1) * sizeof(uint32_t));
        memcpy(right_node->data, node->data + mid_index + 1, (num_count - mid_index - 1) * sizeof(struct info));
        memcpy(right_node->child, node->child + mid_index + 1, (num_count - mid_index) * sizeof(struct node *));

        node->num_keys = mid_index;

        right_node->num_keys = (num_count - mid_index - 1);
        right_node->parent = node->parent;

        node_parent = node->parent;


        //init a node
        if (node_parent == NULL) {
            node_parent = create_tree_node(b_tree);
            b_tree->root = node_parent;
            //init parent node
            node_parent->num_keys++;
            node_parent->keys[0] = node->keys[mid_index];
            node_parent->data[0] = node->data[mid_index];
            node_parent->child[0] = node;
            node_parent->child[1] = right_node;
            right_node->parent = node_parent;
            node->parent = node_parent;
        } else {
            for (index = node_parent->num_keys; index > 0; index--) {
                if (node->keys[mid_index] < node_parent->keys[index - 1]) {
                    node_parent->keys[index] = node_parent->keys[index - 1];
                    node_parent->data[index] = node_parent->data[index - 1];
                    node_parent->child[index + 1] = node_parent->child[index];
                    continue;
                }
                break;
            }
            node_parent->num_keys++;
            right_node->parent = node_parent;
            node_parent->keys[index] = node->keys[mid_index];
            node_parent->data[index] = node->data[mid_index];
            node_parent->child[index + 1] = right_node;

        }

        memset(node->child + mid_index + 1, 0, sizeof(struct node *) * (num_count - mid_index));
        memset(node->keys + mid_index, 0, sizeof(uint32_t) * (num_count - mid_index));
        memset(node->data + mid_index, 0, sizeof(struct info) * (num_count - mid_index));

        for (index = 0; index <= right_node->num_keys; index++) {
            if (right_node->child[index] != NULL) {
                right_node->child[index]->parent = right_node;
            }
        }

        node = node_parent;
    }
    pthread_mutex_unlock(&lock_mutex);
    return 0;
}

int btree_retrieve(uint32_t key, struct info *found, void *helper) {
    pthread_mutex_lock(&lock_mutex);
    struct b_tree *b_tree = helper;
    struct node *node = b_tree->root;
    uint16_t index = 0;
    while(node != NULL && node->num_keys != 0) {
        for (index = 0; index < node->num_keys; index++) {
            //key found in b tree
            if (node->keys[index] == key) {
                memcpy(found->key, (node->data + index)->key, sizeof (uint32_t) * 4);
                found->nonce = (node->data + index)->nonce;
                found->size = (node->data + index)->size;
                found->data = (node->data + index)->data;
                pthread_mutex_unlock(&lock_mutex);
                return 0;
            }
            if (node->keys[index] > key) {
                break;
            }
        }

        //search child
        if (node->child[index] != NULL) {
            node = node->child[index];
        } else {
            break;
        }
    }
    pthread_mutex_unlock(&lock_mutex);
    return 1;
}

int btree_decrypt(uint32_t key, void *output, void *helper) {
    struct info* found = malloc(sizeof (struct info));
    int res = btree_retrieve(key, found, helper);
    if(res == 0){
        uint32_t size = found->size;
        uint64_t nonce = found->nonce;
        void * cipher = found->data;
        uint32_t encryption_key[4];
        memcpy(encryption_key, found->key, sizeof (uint32_t) * 4);
        uint32_t num_blocks = size % 8 == 0 ? size / 8 : (size / 8 + 1);
        uint64_t *plain = (uint64_t *)malloc(sizeof (uint64_t) * num_blocks);
        memset(plain, 0, sizeof (uint64_t) * num_blocks);
        decrypt_tea_ctr(cipher, encryption_key, nonce, plain, num_blocks);
        memcpy(output, plain, size);
        if(plain){
            free(plain);
            plain = NULL;
        }
    }
    return res;
}

int merge_node(struct node * left_brother_node, struct node * right_brother_node, uint32_t idx, struct b_tree* b_tree){
    left_brother_node->keys[left_brother_node->num_keys] = left_brother_node->parent->keys[idx];
    left_brother_node->data[left_brother_node->num_keys] = left_brother_node->parent->data[idx];
    left_brother_node->num_keys++;

    memcpy(left_brother_node->child + left_brother_node->num_keys, right_brother_node->child, sizeof(struct node *) * (right_brother_node->num_keys + 1));
    memcpy(left_brother_node->keys + left_brother_node->num_keys, right_brother_node->keys, sizeof(uint32_t) * (right_brother_node->num_keys));
    memcpy(left_brother_node->data + left_brother_node->num_keys, right_brother_node->data, sizeof(struct info) * (right_brother_node->num_keys));

    for(uint16_t i = 0; i < right_brother_node->num_keys + 1; i++){
        if(right_brother_node->child[i] != NULL){
            right_brother_node->child[i]->parent = left_brother_node;
        }
    }
    left_brother_node->num_keys =  left_brother_node->num_keys + right_brother_node->num_keys;

    for(uint16_t i = idx + 1; i < left_brother_node->parent->num_keys; i++){
        left_brother_node->parent->child[i] = left_brother_node->parent->child[i + 1];
        left_brother_node->parent->keys[i - 1] = left_brother_node->parent->keys[i];
        left_brother_node->parent->data[i - 1] = left_brother_node->parent->data[i];
    }

    left_brother_node->parent->num_keys--;
    left_brother_node->parent->child[left_brother_node->parent->num_keys + 1] = NULL;
    left_brother_node->parent->keys[left_brother_node->parent->num_keys] = 0;
    free(right_brother_node);
   /* if(right_brother_node){
     *//*   for(int i = 0 ; i< right_brother_node->num_keys; i++){
            (right_brother_node->data + i)->data = NULL;
        }*//*
       // free(right_brother_node);
        right_brother_node = NULL;
    }*/


    if(b_tree->min_val > left_brother_node->parent->num_keys){
        return borrowKey(left_brother_node->parent, b_tree);
    }

    return 0;
}

int  borrowKey(struct node * node, struct b_tree *b_tree){
    //root node
    if(node->parent == NULL){
        if(node->num_keys == 0){
            if(node->child[0] == NULL){
                b_tree->root = NULL;
            }else{
                b_tree->root = node->child[0];
                node->child[0]->parent = NULL;
            }
            free(node);
            node = NULL;
        }
        return 0;
    }
    uint16_t index = 0;
    for(index = 0; index < node->parent->num_keys + 1; index++){
        if(node == node->parent->child[index]){
            break;
        }
    }
    //node not found
    if(index == node->parent->num_keys + 1){
        return 1;
    }
    //right brother node
    struct node * left_brother_node = NULL;
    if(index > 0 ){
        left_brother_node = node->parent->child[index -1];
    }
    if(left_brother_node != NULL &&  b_tree->min_val < left_brother_node->num_keys){
        for(uint16_t i = 0; i < node->num_keys; i++){
            node->child[i+1] = node->child[i];
            node->keys[i] = node->keys[i+1];
            node->data[i] = node->data[i+1];
        }

        node->child[1] = node->child[0];

        node->num_keys++;
        node->keys[0] = node->parent->keys[index - 1];
        node->data[0] = node->parent->data[index - 1];

        node->child[0] = left_brother_node->child[left_brother_node->num_keys];
        if(node->child[0] != NULL) {
            node->child[0]->parent = node;
        }

        left_brother_node->num_keys--;
        left_brother_node->child[left_brother_node->num_keys] = NULL;
        node->parent->keys[index - 1] = left_brother_node->keys[left_brother_node->num_keys];
        node->parent->data[index - 1] = left_brother_node->data[left_brother_node->num_keys];
        return 0;
    }
    //right brother node
    struct node * right_brother_node = NULL;
    if(index < node->parent->num_keys ){
        right_brother_node = node->parent->child[index + 1];
    }

    if(right_brother_node != NULL && b_tree->min_val < right_brother_node->num_keys){

        node->keys[node->num_keys] = node->parent->keys[index];
        node->data[node->num_keys] = node->parent->data[index];
        node->num_keys++;

        node->child[node->num_keys] = right_brother_node->child[0];
        if( node->child[node->num_keys] != NULL){
            node->child[node->num_keys]->parent = node;
        }

        node->parent->keys[index] = right_brother_node->keys[0];
        node->parent->data[index] = right_brother_node->data[0];

        for(uint64_t i = 1; i < right_brother_node->num_keys; i++){
            right_brother_node->child[i-1] = right_brother_node->child[i];
            right_brother_node->keys[i-1] = right_brother_node->keys[i];
            right_brother_node->data[i-1] = right_brother_node->data[i];
        }
        right_brother_node->child[right_brother_node->num_keys - 1] = right_brother_node->child[right_brother_node->num_keys];

        right_brother_node->num_keys--;
        right_brother_node->child[right_brother_node->num_keys] = NULL;
        right_brother_node->keys[right_brother_node->num_keys] = 0;
        return 0;
    }else{
        //merge
        uint32_t idx = 0;
        struct node * left_brother_node = NULL;
        struct node * right_brother_node = NULL;
        if(index <= 0){
            left_brother_node = node;
            right_brother_node = node->parent->child[index + 1];
            idx = index;
        }else{
            left_brother_node = node->parent->child[index -1];
            right_brother_node = node;
            idx = index - 1;
        }
        return merge_node(left_brother_node, right_brother_node, idx, b_tree);

    }
}


int btree_delete(uint32_t key, void *helper) {
    pthread_mutex_lock(&lock_mutex);
    struct b_tree *b_tree = helper;
    struct node *node = b_tree->root;
    uint16_t index = 0;
    while(node != NULL && node->num_keys != 0) {
        for (index = 0; index < node->num_keys; index++) {
            //key found in b tree
            if (node->keys[index] == key) {
                if(node->child[index] == NULL){
                    //remove directly
                    for(uint16_t i = index; i < node->num_keys - 1; i++){
                        node->keys[i] = node->keys[i + 1];
                        node->data[i] = node->data[i + 1];
                    }
                    node->num_keys--;
                    node->keys[node->num_keys] = 0;
                    //need free data
                    void *p = (node->data + node->num_keys)->data;
                    (node->data + node->num_keys)->data = NULL;
                    free(p);

                }else{
                    struct node *child_node = node->child[index];
                    while(child_node != NULL && child_node->child[child_node->num_keys] != NULL){
                        child_node = child_node->child[child_node->num_keys];
                    }

                    child_node->num_keys--;
                    node->keys[index] = child_node->keys[child_node->num_keys];
                    void *p = (node->data + index)->data;
                    (node->data + index)->data = NULL;
                    free(p);
                    node->data[index] = child_node->data[child_node->num_keys];

                    child_node->keys[child_node->num_keys] = 0;
                    //need free data
                    node = child_node;

                }

                if(b_tree->min_val > node->num_keys){
                    borrowKey(node, b_tree);
                }
                pthread_mutex_unlock(&lock_mutex);
                return 0;
            }
            if (node->keys[index] > key) {
                break;
            }
        }

        //search child
        if (node->child[index] != NULL) {
            node = node->child[index];
        } else {
            break;
        }
    }
    pthread_mutex_unlock(&lock_mutex);
    return 0;
}

void count_tree_node(struct node *root, uint64_t *pos){
    if(root == NULL){
        return;
    }
    (*pos)++;
    for (uint16_t i = 0; i <= root->num_keys; i++) {
        count_tree_node(root->child[i], pos);
    }
}

void visit_node(struct node *root, uint64_t *pos, struct node **list){
    if(root == NULL){
        return;
    }
    list[*pos] = root;
    (*pos)++;
    for (uint16_t i = 0; i <= root->num_keys; i++) {
        visit_node(root->child[i], pos, list);
    }

}

uint64_t btree_export(void *helper, struct node **list) {
    pthread_mutex_lock(&lock_mutex);
    struct b_tree * tree = helper;
    uint64_t pos = 0;
    visit_node(tree->root, &pos, list);
    pthread_mutex_unlock(&lock_mutex);
    return pos;
}

void encrypt_tea(uint32_t plain[2], uint32_t cipher[2], uint32_t key[4]) {
    uint32_t sum = 0;
    uint32_t delta = 0x9E3779B9;
    cipher[0] = plain[0];
    cipher[1] = plain[1];
    for(uint32_t i = 0; i< 1024; i++){
        sum = (sum + delta) & int_max;
        uint32_t tmp1 = ((cipher[1] << 4) + key[0]) & int_max;
        uint32_t tmp2 = (cipher[1] + sum) & int_max;
        uint32_t tmp3 = ((cipher[1] >> 5) + key[1]) & int_max;
        cipher[0] = (cipher[0] + (tmp1 ^ tmp2 ^ tmp3)) & int_max;
        uint32_t tmp4 = ((cipher[0] << 4) + key[2]) & int_max;
        uint32_t tmp5 = (cipher[0] + sum) & int_max;
        uint32_t tmp6 = ((cipher[0] >> 5) + key[3]) & int_max;
        cipher[1] = (cipher[1] + (tmp4 ^ tmp5 ^ tmp6)) & int_max;
    }
}

void decrypt_tea(uint32_t cipher[2], uint32_t plain[2], uint32_t key[4]) {
    uint32_t sum = 0xDDE6E400;
    uint32_t delta = 0x9E3779B9;
    for(uint32_t i = 0; i< 1024; i++){
        uint32_t tmp4 = ((cipher[0] << 4) + key[2]) & int_max;
        uint32_t tmp5 = (cipher[0] + sum) & int_max;
        uint32_t tmp6 = ((cipher[0] >> 5) + key[3]) & int_max;
        cipher[1] = (cipher[1] - (tmp4 ^ tmp5 ^ tmp6)) & int_max;
        uint32_t tmp1 = ((cipher[1] << 4) + key[0]) & int_max;
        uint32_t tmp2 = (cipher[1] + sum) & int_max;
        uint32_t tmp3 = ((cipher[1] >> 5) + key[1]) & int_max;
        cipher[0] = (cipher[0] - (tmp1 ^ tmp2 ^ tmp3)) & int_max;
        sum = (sum - delta) & int_max;
    }
    plain[0] = cipher[0];
    plain[1] = cipher[1];
}

void encrypt_tea_ctr(uint64_t * plain, uint32_t key[4], uint64_t nonce, uint64_t * cipher, uint32_t num_blocks) {
    for(uint32_t i = 0; i < num_blocks; i++){
        uint64_t tmp1 = i ^ nonce;
        uint32_t c[2];
        uint32_t p[2];
        p[1]  = tmp1 >> 32;
        p[0] = tmp1 & int_max;
        encrypt_tea(p, c, key);
        uint64_t tmp3 = c[1];
        uint64_t block_cipher =  (tmp3 << 32) + c[0];
        cipher[i] = plain[i] ^ (block_cipher);
    }
}

void decrypt_tea_ctr(uint64_t * cipher, uint32_t key[4], uint64_t nonce, uint64_t * plain, uint32_t num_blocks) {

    for(uint32_t i = 0; i < num_blocks; i++){
        uint64_t tmp1 = i ^ nonce;
        uint32_t c[2];
        uint32_t p[2];
        p[1]  = tmp1 >> 32;
        p[0] = tmp1 & int_max;
        encrypt_tea(p, c, key);
        uint64_t tmp3 = c[1];
        uint64_t block_cipher =  (tmp3 << 32) + c[0];
        plain[i] = cipher[i] ^ (block_cipher);
    }
}